;;;; -*- Mode: Lisp; Package: :lib.gl; -*-

;;;; ---- ;;;;
;;; Supporting array and vectors that can be used in GL and in Lisp.
;;;

(in-package :lib.gl)

(defparameter *external-int-size* (ccl::foreign-size :<GL>int :bytes))
(defparameter *external-double-size* (ccl::foreign-size :<GL>double :bytes))
(defparameter *external-float-size* (ccl::foreign-size :<GL>float :bytes))

;;; ------------------------------------------------------------
;;; Types and creators.
(defclass external-pointer ()
  ((element-type :reader element-type :initarg :element-type)
   (element-count :reader element-count :initarg :element-count)
   (dims :reader dims :initarg :dims)
   (ivector :reader ivector :initarg :ivector)
   (macptr :reader macptr)))

(defmethod xml::xml-printable-as-attribute-value-p ((obj external-pointer))
  t)

(defmethod xml::print-typed-attribute-value (value (type (eql 'external-pointer)) stream)
  (xml::print-typed-attribute-value (ivector value) 'array stream))

(defmethod get-pointer ((self external-pointer))
  (slot-value self 'macptr))

(defmethod pointer-size ((self external-pointer))
  (with-slots (element-type element-count) self
    (ecase element-type
      (double-float (* *external-double-size* element-count))
      (single-float (* *external-float-size* element-count))
      (fixnum (* *external-int-size* element-count))
      (t (* *external-int-size* element-count)))))

(defmethod external-array-dimension ((self external-pointer) i)
  (with-slots (dims) self
    (nth i dims)))

(defmethod initialize-instance :around ((self external-pointer) &rest initargs)
  (declare (ignore initargs))
  (call-next-method)
  (ccl:terminate-when-unreachable self)
  (with-slots (ivector macptr element-type element-count) self
    (multiple-value-bind (new-ivector new-macptr)
	(ccl:make-heap-ivector element-count element-type)
      (setq ivector new-ivector
	    macptr new-macptr))))


(defmethod ccl:terminate ((self external-pointer))
  (with-slots (ivector macptr) self
    (when ivector
      (ccl:dispose-heap-ivector ivector)
      (setq ivector nil
	    macptr nil))))

;; Shamelessly stolen from CCL's innards.
(defun calc-index (rank dims subscripts)
  (let ((sum 0))
    (do* ((i (1- rank) (1- i))
	  (dim (nth i dims) (nth i dims))
	  (last-size 1 size)
	  (size dim (* dim size)))
	 (nil)
      (declare (fixnum i last-size size))
      (let ((s (nth i subscripts)))
	(if (>= s dim)
	    (error "Subscript ~A out of range for dimension of ~A." s dim))
	(unless (ccl:fixnump s)
	  (setq s (ccl:require-type s 'fixnum)))
	(incf sum (the fixnum (* s last-size)))
	(when (eql i 0) (return sum))))))

(defmacro safe-external-assign (place op)
  `(let ((fac ,op))
     (restart-case
	 (handler-case (setf ,place fac)
	   (ccl::array-element-type-error () (invoke-restart 'safe-assign)))
       (safe-assign () (let ((etype (type-of ,place)))
		      (setf ,place (coerce fac etype)))))))

(defmethod external-aref ((self external-pointer) &rest idx)
  (with-slots (ivector dims) self
    (if (> (length idx) 1)
	(aref ivector (calc-index (length dims) dims idx)) ;(setq idx (list (calc-index (length dims) dims idx))))
	(aref ivector (car idx)))));; (apply #'aref ivector idx)

(defmethod (setf external-aref) (val (self external-pointer) &rest idx)
  (with-slots (ivector dims) self
    (if (> (length idx) 1)
	(progn
	  (setq idx (calc-index (length dims) dims idx));(list (calc-index (length dims) dims idx))))
	  (safe-external-assign (aref ivector idx) val))
	(safe-external-assign (aref ivector (car idx)) val))))

(defclass external-vector (external-pointer)
  ())

(defmethod xml::xml-printable-as-attribute-value-p ((obj external-vector))
  t)

(defmethod xml::print-typed-attribute-value (value (type (eql 'external-vector)) stream)
  (xml::print-typed-attribute-value (ivector value) 'array stream))

(defmethod make-load-form ((obj external-vector) &optional env)
  (declare (ignore env))
  (with-slots (ivector element-type dims) obj
    `(make-external-array ,dims :initial-contents ',(coerce ivector 'list) :element-type ',element-type)))

(defun print-vector (vector &optional (stream *standard-output*))
  (with-slots (ivector) vector
    (princ ivector stream)))

(defmethod print-object ((vec external-vector) stream)
  (print-unreadable-object (vec stream :type t)
    (print-vector vec stream)))

(defun %external-vector (&rest elements)
  (make-external-array (length elements) 
		 :initial-contents elements 
		 :element-type (typecase (first elements)
				 (integer 'fixnum)
				 (single-float 'single-float)
				 (double-float 'double-float)
				 (t 'fixnum))))
					;    (make-external-array (length elements) :initial-contents elements :element-type 'fixnum)
 ;   (make-external-array (length elements) :initial-contents elements :element-type (type-of (first elements)))))

(iterate:defmacro-clause (FOR var IN-EXTERNAL-VECTOR v)
  (alexandria:with-unique-names (index vect end)
      `(progn
	 (with ,vect = ,v)
	 (with ,end = (slot-value ,vect 'element-count))

	 (with ,index = -1)
	 (for ,var next (progn (incf ,index)
			       (if (>= ,index ,end) (finish))
			       (external-aref ,vect ,index))))))

	    
(defun external-vector-length (external-vector)
  "Return the number of elements in the foreign vector."
  (with-slots (ivector) external-vector
    (length ivector)))

(defun external-vector-size (evec)
  "Return the size, in bytes, of the foreign vector."
  (* (external-vector-length evec) (pointer-size evec)))

(defclass external-array (external-pointer)
  ())

(defmethod xml::xml-printable-as-attribute-value-p ((obj external-array))
  t)

(defmethod xml::print-typed-attribute-value (value (type (eql 'external-array)) stream)
  (xml::print-typed-attribute-value (ivector value) 'array stream))

;;; Maybe this could print prettier, but it will
;;; have to serve.
(defmethod print-object ((vec external-array) stream)
  (print-unreadable-object (vec stream :type t :identity t)
    (print-vector vec stream)))

(defun make-external-array (dims &key (initial-contents nil) (element-type 'single-float)
		      (initial-element nil))
  (if (and initial-element (not (eql element-type (type-of initial-element))))
    (setq element-type (type-of initial-element)))
  (labels ((find-element-type (nested-list)
	     (if (listp (car nested-list))
		 (find-element-type (car nested-list))
		 (type-of (car nested-list)))))
    (if (and initial-contents (not element-type))
	(setq element-type (find-element-type initial-contents))))
  (let ((vec (if (listp dims)
		 (make-instance 'external-array :element-type element-type
				:dims dims
				:element-count (if (listp dims)
						   (apply #'* dims)
						   dims))
		 (make-instance 'external-vector :element-type element-type
				:dims dims
				:element-count (if (listp dims)
						   (apply #'* dims)
						   dims)))))
    (if initial-element
	(iterate (for i from 0 to (1- (if (listp dims)
					  (apply #'* dims)
					  dims)))
		 (after-each (setf (external-aref vec i) initial-element)))
	(let ((cnt 0))
	  (flet ((set-ref (vec val idx)
		   (with-slots (ivector) vec
		     (setf (aref ivector idx) val))))
	    (when initial-contents
	      (if (listp (car initial-contents))
		  (iter (for line in initial-contents)
			(after-each (iter (for elem in line)
					  (after-each (set-ref vec elem cnt)
						      (incf cnt)))))
		  (iter (for elem in initial-contents)
			(after-each (set-ref vec elem cnt)
				    (incf cnt))))))))
    vec))

(defun external-array-equal (a1 a2)
  (if (not (= (element-count a1) (element-count a2)))
    nil)
  (dotimes (i (element-count a1) t)
    (when (not (= (external-aref a1 i) (external-aref a2 i)))
      (return nil))))

(defun array->external-array (a &optional (element-type nil))
  (let* ((dims (array-dimensions a))
         (size (array-total-size a))
         (gla (make-external-array dims :element-type (or element-type (type-of (row-major-aref a 0))))))
    (dotimes (i size gla)
      (setf (external-aref gla i) (row-major-aref a i)))))

;;; Define #{list} to be a call to (%external-vector list)

(defun |#{-reader| (stream char arg)
  (declare (ignore char arg))
  (apply #'%external-vector (read-delimited-list #\} stream t)))

(defparameter *gltypes-readtable* (copy-readtable))

(let ((*readtable* *gltypes-readtable*))
  (set-dispatch-macro-character #\# #\{ #'|#{-reader|)
  (set-macro-character #\} (get-macro-character #\) nil)))

(defun enable-external-type-syntax ()
  (setq *readtable* *gltypes-readtable*))

(let ((clean-readtable (copy-readtable)))
  (defun disable-external-type-syntax ()
    (setq *readtable* clean-readtable)))

;;; Silly Utilities

(defun swap-elements (glvec s i)
  (psetf (external-aref glvec s) (external-aref glvec i)
	 (external-aref glvec i) (external-aref glvec s)))

(defun fisher-yates-shuffle (glvec)
  "Destructively modifies the contents of glvec, randomly ordering its contents using the fisher-yates algorithm."
  (when (> (element-count glvec) 1)
    (iter (for i from (1- (element-count glvec)) downto 1)
	  (let ((s (random i)))
	    (swap-elements glvec s i))))
  glvec)

(defun test-shuffle ()
  (format t "~A~%" (fisher-yates-shuffle (make-vector 1 2 3 4 5 6 7 8 9 10))))

(defun test-make-array ()
  (make-external-array 3 :initial-element 1.0))

(defun make-external-index-array (&rest indices)
  (make-external-array (length indices)
		 :initial-contents indices
		 :element-type '(unsigned-byte 32)))

;;; Working with cl-opengl

;; The cl-opengl version only returns a pointer, I just wanted to save
;; myself the work of creating gl-array structs...
(defmacro with-opengl-array ((var type lisp-array) &body body)
  (check-type var symbol)
  (let ((count (gensym "COUNT"))
        (array (gensym "ARRAY"))
        (original-type type))
    (alexandria:once-only (type)
      `(let* ((,array ,lisp-array)
              (,count (length ,array)))
         (cffi:with-foreign-object (,var ,type ,count)
           ;; we need type to be a constant within the loop so cffi can
           ;; optimize it, so check type outside the loop, and make a
           ;; copy of the loop for any types we care about
           ;; (unless we know the type at compile time, in which case only
           ;;  generate code for that type)
           ,(if (or (keywordp original-type)
                    (and (consp original-type)
                         (eq (first original-type) 'quote)))
                `(loop for i below ,count
                    do (setf (cffi:mem-aref ,var ,original-type i)
                             (aref ,array i)))
                `(case ,type
                   ,@(loop for ctype in '(%gl:byte %gl:ubyte %gl:short
                                          %gl:ushort %gl:int %gl:uint %gl:float)
                        collect
                          `(,ctype
                            (loop for i below ,count
                               do (setf (mem-aref ,var ',ctype i) (aref ,array i)))))
                   (t
                    (loop for i below ,count
                       do (setf (mem-aref ,var ,type i) (aref ,array i))))))
	   (let ((,var (gl::make-gl-array-from-pointer ,var ,type ,count)))
	     ,@body))))))

;;; Matrices and Vectors

(defun print-external-matrix (matrix &optional (stream *standard-output*))
  (if (column-vector-p matrix)
      (iter (for x in-external-vector matrix)
	    (format stream "| ~9,2F |~%" x))
      (dotimes (r (matrix-rows matrix))
	(format stream "| ")
	(dotimes (c (matrix-cols matrix))
	  (format stream "~9,2F" (mref matrix r c)))
	(format stream " |")
	(terpri stream))))

(defmethod print-matrix ((m external-array) &rest args)
  (apply #'print-external-matrix m args))

(defmethod matrix-rows ((matrix external-array))
  (external-array-dimension matrix 0))

(defmethod matrix-cols ((matrix external-array))
  (external-array-dimension matrix 1))

(defmethod mref ((matrix external-array) row &optional col)
  (if (null col)
      (external-aref matrix row)
      (external-aref matrix col row)))

(defmethod (setf mref) (val (matrix external-array) &rest indices)
  (setf (external-aref matrix (second indices) (first indices))
        val))

(defmethod row-major-mref ((matrix external-array) idx)
  (external-aref matrix idx))

(defmethod (setf row-major-mref) (val (matrix external-array) idx)
  (setf (external-aref matrix idx) val))

(defmethod m/ ((matrix external-array) (k number) &rest args)
  (declare (ignore args))
  (let ((r (matrix-rows matrix)))
    (if (not (= k 0))
        (let ((kinv (/ +one+ k)))
          (dotimes (i (* r r) matrix)
            (setf (external-aref matrix i)
                  (* (external-aref matrix i) kinv))))
	(dotimes (i (* r r) matrix)
	  (setf (external-aref matrix i) most-positive-double-float)))))

(defmethod m* ((matrix external-array) (k number) &rest args)
  (apply #'m* k matrix args))

(defmethod m* ((k number) (matrix external-array) &rest args)
  (let ((r (%mul-matrices-scalar k matrix)))
    (if (not (null args))
        (apply #'m* r (first args) (rest args))
	r)))

(defmethod m* ((p external-array) (q external-array) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-matrices p q) (first rest) (rest rest))
    (%mult-matrices p q)))

(defmethod m* ((p external-array) (q external-vector) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-matvec p q) (first rest) (rest rest))
      (%mult-matvec p q)))

(defmethod m* ((q external-vector) (p external-array) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-vecmat p q) (first rest) (rest rest))
      (%mult-vecmat q p)))

(defmethod m- ((m1 external-array) (m2 external-array) &rest args)
  (if (not (null args))
      (apply #'m- (%sub-mat m1 m2) (first args) (rest args))
    (%sub-mat m1 m2)))

(defmethod mneg ((m external-array))
  (dotimes (i (* (matrix-rows m)
                 (matrix-cols m))
              m)
    (setf (row-major-mref m i)
          (- (row-major-mref m i)))))

(defmethod m+ ((m1 external-array) (m2 external-array) &rest args)
  (if (not (null args))
      (apply #'m+ (%add-mat m1 m2) (first args) (rest args))
    (%add-mat m1 m2)))

(defmethod times-diagonal ((m external-array) (diag external-vector))
  (make-square-matrix
   (* (external-aref m 0) (external-aref diag 0))
   (* (external-aref m 1) (external-aref diag 1))
   (* (external-aref m 2) (external-aref diag 2))
   (* (external-aref m 3) (external-aref diag 0))
   (* (external-aref m 4) (external-aref diag 1))
   (* (external-aref m 5) (external-aref diag 2))
   (* (external-aref m 6) (external-aref diag 0))
   (* (external-aref m 7) (external-aref diag 1))
   (* (external-aref m 8) (external-aref diag 2))))

(defmethod diagonal-times ((m external-array) (diag external-vector))
  (make-square-matrix
   (* (external-aref m 0) (external-aref diag 0))
   (* (external-aref m 1) (external-aref diag 0))
   (* (external-aref m 2) (external-aref diag 0))
   (* (external-aref m 3) (external-aref diag 1))
   (* (external-aref m 4) (external-aref diag 1))
   (* (external-aref m 5) (external-aref diag 1))
   (* (external-aref m 6) (external-aref diag 2))
   (* (external-aref m 7) (external-aref diag 2))
   (* (external-aref m 8) (external-aref diag 2))))

(defmethod mref ((v external-vector) row &optional col)
  (declare (ignore col))
  (external-aref v row))

(defmethod (setf mref) (val (v external-vector) &rest indices)
  (setf (external-aref v (car indices)) val))

(defmethod vector-length ((vec external-vector))
  (element-count vec))

(defmethod copy-vector ((v external-vector))
  (let ((vn (zero-vector (vector-length v))))
    (iter (for vi in-external-vector v)
	  (for i from 0)
	  (setf (mref vn i) vi))
    vn))

(defmethod m+ ((v1 external-vector) (v2 external-vector) &rest args)
  (if (not (null args))
      (apply #'m+ (%add-vec v1 v2) (first args) (rest args))
    (%add-vec v1 v2)))

(defmethod m- ((v1 external-vector) (v2 external-vector) &rest args)
  (if (not (null args))
      (apply #'m- (%sub-vec v1 v2) (first args) (rest args))
    (%sub-vec v1 v2)))

(Defmethod m* ((k number) (q external-vector) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-scalar-vec k q) (first rest) (rest rest))
    (%mult-scalar-vec k q)))

(defmethod m* ((q external-vector) (u external-vector) &rest rest)
  (if (not (null rest))
      (apply #'m* (%mult-vec-vec q u) (first rest) (rest rest))
    (%mult-vec-vec q u)))

(defmethod m/ ((q external-vector) (k number) &rest rest)
  (if (not (null rest))
      (apply #'m* (%div-scalar-vec k q) (first rest) (rest rest))
    (%div-scalar-vec k q)))

(defmethod mneg ((v external-vector))
  (let* ((l (vector-length v))
	 (nvec (make-external-array l :initial-element (mref v 0))))
    (dotimes (i l nvec)
      (setf (mref nvec i) (- (mref v i))))))

(defmethod magnitude ((v external-vector))
  (sqrt (iter (for n in-external-vector v)
              (summing (* n n)))))

(defmethod squared-magnitude ((v external-vector))
  (iter (for n in-external-vector v)
        (summing (* n n))))

(defmethod dot-product ((v1 external-vector) (v2 external-vector))
  (loop for i upto (1- (external-vector-length v1))
        summing (* (mref v1 i) (mref v2 i))))

